import os
import uuid

from flask import Flask,jsonify,request, make_response, send_file
import logging
from skimage import io
import torch
from PIL import Image
from briarmbg import BriaRMBG
from utilities import preprocess_image, postprocess_image
from io import BytesIO

app = Flask(__name__)

# 创建一个日志器
logger = logging.getLogger(__name__)


def random_string(length):
    result = str(uuid.uuid4())[:length]
    return result

def delete_file(filename):
    if os.path.exists(filename):
        os.remove(filename)
        logger.info(f"文件 {filename} 删除成功！")
    else:
        logger.info(f"文件 {filename} 不存在。")

# 第一个路由 /hello：
# 路径：/hello
# 允许的请求方法：['GET', 'POST']
# 视图函数：hello，当访问 /hello 路径时，无论是 GET 还是 POST 请求，都会调用这个函数，并返回 "Hello World!"。
# methods参数用于指定允许的请求格式
# 常规输入url的访问就是get方法

net = BriaRMBG()
device = torch.device("cpu")
net.load_state_dict(torch.load("./model.pth", map_location=torch.device('cpu')))
net.to(device)

@app.route("/image", methods=['POST'])
def image():
    if 'file' not in request.files:
        resp = make_response('no file', 500)
        return  resp
    file = request.files['file']
    # 如果用户没有选择文件，浏览器也会提交一个没有文件名的空部分
    if file.filename == '':
        resp = make_response('No selected file', 500)
        return resp
    # 如果文件名是空的，那么也不进行保存
    if file:
        # 保存文件到服务器的uploads文件夹，文件名使用secure_filename保证安全
        file_extension = os.path.splitext(file.filename)[1]
        image_file_path = 'imageTemp/' + random_string(8) + file_extension
        file.save(image_file_path)
        # prepare input
        model_input_size = [1024, 1024]
        orig_im = io.imread(image_file_path)
        orig_im_size = orig_im.shape[0:2]
        image = preprocess_image(orig_im, model_input_size).to(device)

        # inference
        result = net(image)

        # post process
        result_image = postprocess_image(result[0][0], orig_im_size)

        # save result
        pil_im = Image.fromarray(result_image)
        no_bg_image = Image.new("RGBA", pil_im.size, (0, 0, 0, 0))
        orig_image = Image.open(image_file_path)
        no_bg_image.paste(orig_image, mask=pil_im)
        # convert_file_path = 'imageTemp/'+random_string(8) + '.webp'
        # no_bg_image.save(convert_file_path)
        delete_file(image_file_path)
        binary_file = BytesIO()
        no_bg_image.save(binary_file, format='WEBP')
        binary_file.seek(0)
        return send_file(
            binary_file,
            mimetype='image/webp',
            as_attachment=True,
            download_name='result.webp'
        )
        # with open(convert_file_path, 'rb') as file:
        #     data = file.read()
        #     base64_str  = base64.b64encode(data).decode('utf-8')
        #     json_res = {
        #         'code': 0,
        #         'data': base64_str,
        #     }
        #     file.close()
        #     delete_file(image_file_path)
        #     delete_file(convert_file_path)
        #     return jsonify(json_res)


# 注意路由路径不要重名，映射的视图函数也不要重名
# 第二个路由 /hi：
# 路径：/hi
# 允许的请求方法：['POST']
# 视图函数：hi，这个路由只响应 POST 请求，当有 POST 请求访问 /hi 路径时，调用 hi 函数，并返回 "Hi World!"。

app.run(host='0.0.0.0')
